"""
This file contains some useful functions for train / val.
"""
import os
import numpy as np
import torch
import collections


#################
## image utils ##
#################
def convert_image(input_tensor, axis):
    """ Convert single channel images to 3-channel images. """
    image_lst = [input_tensor for _ in range(3)]
    outputs = np.concatenate(image_lst, axis)
    return outputs


######################
## checkpoint utils ##
######################
def get_latest_checkpoint(checkpoint_root, checkpoint_name,
                          device=torch.device("cuda")):
    """ Get the latest checkpoint or by filename. """
    # Load specific checkpoint
    if checkpoint_name is not None:
        checkpoint = torch.load(
            os.path.join(checkpoint_root, checkpoint_name),
            map_location=device)
        #去除多gpu模型字典中的'module'
        for k,v in checkpoint.items():
            new_k = k.replace('module.','') if 'module' in k else k
            checkpoint[new_k]=v
            if type(checkpoint[new_k])==dict:
                for key,value in checkpoint[new_k].items():
                    new_key = key.replace('module.', '') if 'module' in key else key
                    checkpoint[new_k][new_key]=v
            if type(checkpoint[new_k])== collections.OrderedDict:
                ck = collections.OrderedDict()
                for key,value in checkpoint[new_k].items():
                    if 'module' in key:
                        new_key=key.replace('module.','')
                        ck[new_key]=value
                    else:
                        ck[key]=value
                checkpoint[new_k]=ck
    # Load the latest checkpoint
    else:
        lastest_checkpoint = sorted(os.listdir(os.path.join(
            checkpoint_root, "*.tar")))[-1]
        checkpoint = torch.load(os.path.join(
            checkpoint_root, lastest_checkpoint), map_location=device)
    return checkpoint


# max_ckpt=150
def remove_old_checkpoints(checkpoint_root, max_ckpt=15):
    """ Remove the outdated checkpoints. """
    # Get sorted list of checkpoints
    checkpoint_list = sorted(
        [_ for _ in os.listdir(os.path.join(checkpoint_root))
         if _.endswith(".tar")])

    # Get the checkpoints to be removed
    if len(checkpoint_list) > max_ckpt:
        remove_list = checkpoint_list[:-max_ckpt]
        for _ in remove_list:
            full_name = os.path.join(checkpoint_root, _)
            # 删除文件
            os.remove(full_name)
            print("[Debug] Remove outdated checkpoint %s" % (full_name))

def adapt_checkpoint(state_dict):
    new_state_dict = {}
    for k, v in state_dict.items():
        if k.startswith('module.'):
            new_state_dict[k[7:]] = v
        else:
            new_state_dict[k] = v
    return new_state_dict

################
## HDF5 utils ##
################
def parse_h5_data(h5_data):
    """ Parse h5 dataset. """
    output_data = {}
    for key in h5_data.keys():
        output_data[key] = np.array(h5_data[key])

    return output_data
